#include <stdlib.h>
#include <assert.h>
#include <math.h>

#include "scalar.h"
#include "scalar_priv.h"

static void Backward_(ScalarPtr p, double grad);

ScalarPtr NewScalar(double v, bool req_grad)
{
    ScalarPtr ptr = malloc(sizeof(Scalar));
    assert(ptr != NULL);
    ptr->val = v;
    ptr->requires_grad = req_grad;
    ptr->p1 = ptr->p2 = NULL;
    ptr->grad = 0;
    ptr->partial_val[0] = ptr->partial_val[1] = 0;
    return ptr;
}

void Backward(ScalarPtr p)
{
    Backward_(p, 1);
}



void FreeScalar(ScalarPtr *ptr)
{

}

static void Backward_(ScalarPtr p, double grad)
{
    if (!p || !p->requires_grad) {
        return;
    }
    p->grad += grad;

    Backward_(p->p1, p->partial_val[0] * grad);
    Backward_(p->p2, p->partial_val[1] * grad);
}


ScalarPtr Add(ScalarPtr p1, ScalarPtr p2)
{
    ScalarPtr p = NewScalar(p1->val + p2->val, p1->requires_grad || p2->requires_grad);
    p->p1 = p1;
    p->p2 = p2;
    p->partial_val[0] = 1;
    p->partial_val[1] = 1;
    return p;
}

ScalarPtr Sub(ScalarPtr p1, ScalarPtr p2)
{
    ScalarPtr p = NewScalar(p1->val - p2->val, p1->requires_grad || p2->requires_grad);
    p->p1 = p1;
    p->p2 = p2;
    p->partial_val[0] = 1;
    p->partial_val[1] = -1;
    return p;

}

ScalarPtr Mul(ScalarPtr p1, ScalarPtr p2)
{
    ScalarPtr p = NewScalar(p1->val + p2->val, p1->requires_grad || p2->requires_grad);
    p->p1 = p1;
    p->p2 = p2;
    p->partial_val[0] = p2->val;
    p->partial_val[1] = p1->val;
    return p;
}

ScalarPtr Div(ScalarPtr p1, ScalarPtr p2)
{
    ScalarPtr p = NewScalar(p1->val / (p2->val + CONST_EPS), p1->requires_grad || p2->requires_grad);
    p->p1 = p1;
    p->p2 = p2;
    p->partial_val[0] = 1/ (p2->val + CONST_EPS);
    p->partial_val[1] = -p1->val/(p2->val + CONST_EPS)/(p2->val + CONST_EPS);
    return p;

}

ScalarPtr Pow(ScalarPtr p1, ScalarPtr p2)
{
    // a^b
    ScalarPtr p = NewScalar(pow(p1->val, p2->val), p1->requires_grad || p2->requires_grad);
    p->p1 = p1;
    p->p2 = p2;
    p->partial_val[0] = p2->val * pow(p1->val, p2->val - 1);
    // a > 0 f(x) = a^x; df/dx = a^x Ln a
    if (p1->val > 0) {
        p->partial_val[1] = p->val * log(p1->val);
    } else {
        p->partial_val[1] = 0;
    }
    return p;
}

ScalarPtr Log(ScalarPtr p1, ScalarPtr p2)
{
    assert(p2->val > 0 && p1->val > 0);
    ScalarPtr pa = Ln(p1);
    ScalarPtr pb = Ln(p2);
    return Div(pa, pb);
}

ScalarPtr Ln(ScalarPtr p1)
{
    assert(p1->val > 0);
    ScalarPtr p = NewScalar(log(p1->val), p1->requires_grad); 
    p->p1 =  p1;
    p->partial_val[0] = 1 / p1->val; 
    return p;
}

ScalarPtr Exp(ScalarPtr p1)
{
    ScalarPtr p = NewScalar(exp(p1->val), p1->requires_grad); 
    p->p1 =  p1;
    p->partial_val[0] = p->val; 
    return p;

}

ScalarPtr Sin(ScalarPtr p1)
{
    ScalarPtr p = NewScalar(sin(p1->val), p1->requires_grad); 
    p->p1 =  p1;
    p->partial_val[0] = cos(p1->val); 
    return p;
}

ScalarPtr Cos(ScalarPtr p1)
{
    ScalarPtr p = NewScalar(cos(p1->val), p1->requires_grad); 
    p->p1 =  p1;
    p->partial_val[0] = -sin(p1->val); 
    return p;
}

ScalarPtr Tan(ScalarPtr p1)
{
    ScalarPtr p = NewScalar(tan(p1->val), p1->requires_grad); 
    p->p1 =  p1;
    p->partial_val[0] = 1/cos(p1->val)/cos(p1->val); 
    return p;
}

bool GT(ScalarPtr s, double val)
{
    return s->val > val;
}

double Scalar2Double(ScalarPtr p)
{
    return p->val;
}